package com.yeziji.excel;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.EasyExcelFactory;
import com.alibaba.excel.context.AnalysisContext;
import com.alibaba.excel.event.AnalysisEventListener;
import com.alibaba.excel.read.builder.ExcelReaderBuilder;
import com.alibaba.excel.read.metadata.ReadSheet;
import com.alibaba.excel.support.ExcelTypeEnum;
import com.google.common.collect.Lists;
import com.yeziji.annotation.ExcelHeader;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.poi.poifs.filesystem.FileMagic;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.function.Predicate;

/**
 * easy excel 分批执行器
 *
 * @author hwy
 * @since 2023/06/06 14:26
 **/
@Slf4j
public class EasyExcelBatchExecutor<T> extends AnalysisEventListener<T> {
    /**
     * 默认批量处理的阈值
     */
    public static final int DEFAULT_EXCEL_BATCH = 20000;
    /**
     * 处理的列表
     */
    private List<T> list;
    /**
     * 是否全读
     */
    private boolean readAll;
    /**
     * 起始读的行，一般情况下就一行表头，所以这里默认跳过，如果要设置读头部进 list 就要手动设置为 0
     */
    private int headerRowNumber = 1;
    /**
     * 对应的读取类
     */
    private Class<T> clazz;
    /**
     * 自定义批量
     */
    private int batchSize;
    /**
     * 用户自定义每次批量的操作消费行为
     */
    private BiConsumer<List<T>, Integer> listConsumer;
    /**
     * 用户自定义每次批量携带存储的操作行为
     */
    private BiConsumer<List<T>, Map<T, String>> listValidateConsumer;
    /**
     * 用户自定义校验消费行为
     */
    private Function<T, String> validateFunction;
    /**
     * 失败明细返回原来的值
     */
    private Map<T, String> failedMap;
    /**
     * 是否失败继续进行
     */
    private boolean failedContinue = true;
    /**
     * 约束excel头部，约束之后会自动校验头部
     */
    private List<String> excelHeaders;
    /**
     * 用户自定义头部消费行为
     */
    private Predicate<Map<Integer, String>> headerPredicate;
    /**
     * 头部校验失败的异常信息
     */
    private String excelHeaderErrorMsg = "导入的 excel 模板不正确";
    /**
     * 是否启用清理行为，启用后每次分批执行都会对 list 进行清空
     */
    private boolean clear;
    /**
     * sheet 页码
     */
    private int sheetNo = 0;
    /**
     * 当前页码
     */
    private int currentNo = 0;
    /**
     * 逻辑分批计数器
     */
    private int logicBatch = 0;
    /**
     * 是否清除list开启下一页
     */
    private boolean clearNextSheet = true;
    /**
     * 多页读取的页码列表
     */
    private final List<ReadSheet> sheetNoList = new ArrayList<>();
    /**
     * 数据总量
     */
    private int total = 0;
    /**
     * 表头计数器
     */
    private int recordHeaderCount = 0;
    /**
     * 命中检查的表头
     */
    private Integer hitCheckHeaderNum;
    /**
     * 设置最大数量
     */
    private int maxTotal;
    /**
     * 溢出最大数量时的信息
     */
    private String overflowTotalMsg;

    public EasyExcelBatchExecutor(int batchSize, boolean clear, Class<T> clazz) {
        this.batchSize = batchSize;
        this.clear = clear;
        this.clazz = clazz;
        this.list = new LinkedList<>();
    }

    public static <T> EasyExcelBatchExecutor<T> builder(final Class<T> clazz) {
        EasyExcelBatchExecutor<T> executor = new EasyExcelBatchExecutor<>(DEFAULT_EXCEL_BATCH, false, clazz);
        executor.clazz = clazz;
        executor.batchSize = DEFAULT_EXCEL_BATCH;
        // 默认空构造
        executor.listConsumer = (list, curNo) -> {};
        return executor;
    }

    /* ---------- 支持 buildMap ---------- */
    public static EasyExcelBatchExecutor<Map<Integer, String>> buildMap(final int batchSize) {
        EasyExcelBatchExecutor<Map<Integer, String>> executor = new EasyExcelBatchExecutor<>(batchSize, false, null);
        executor.batchSize = batchSize;
        // 默认空构造
        executor.listConsumer = (list, curNo) -> {};
        return executor;
    }

    public static EasyExcelBatchExecutor<Map<Integer, String>> buildMap(final BiConsumer<List<Map<Integer, String>>, Integer> listConsumer) {
        EasyExcelBatchExecutor<Map<Integer, String>> executor = new EasyExcelBatchExecutor<>(DEFAULT_EXCEL_BATCH, false, null);
        executor.listConsumer = listConsumer;
        return executor;
    }

    public static EasyExcelBatchExecutor<Map<Integer, String>> buildMap(final int batchSize, final BiConsumer<List<Map<Integer, String>>, Integer> listConsumer) {
        EasyExcelBatchExecutor<Map<Integer, String>> executor = new EasyExcelBatchExecutor<>(batchSize, false, null);
        executor.batchSize = batchSize;
        executor.listConsumer = listConsumer;
        return executor;
    }
    /* ---------- buildMap end ---------- */

    /* ---------- 支持自定义 class 映射 ---------- */
    public static <T> EasyExcelBatchExecutor<T> builder(final int batchSize, final Class<T> clazz) {
        EasyExcelBatchExecutor<T> executor = new EasyExcelBatchExecutor<>(batchSize, false, clazz);
        executor.clazz = clazz;
        executor.batchSize = batchSize;
        // 默认空构造
        executor.listConsumer = (list, curNo) -> {};
        return executor;
    }

    public static <T> EasyExcelBatchExecutor<T> builder(final Class<T> clazz, final BiConsumer<List<T>, Integer> listConsumer) {
        EasyExcelBatchExecutor<T> executor = new EasyExcelBatchExecutor<>(DEFAULT_EXCEL_BATCH, false, clazz);
        executor.clazz = clazz;
        executor.listConsumer = listConsumer;
        return executor;
    }

    public static <T> EasyExcelBatchExecutor<T> builder(final int batchSize, final Class<T> clazz, final BiConsumer<List<T>, Integer> listConsumer) {
        EasyExcelBatchExecutor<T> executor = new EasyExcelBatchExecutor<>(batchSize, false, clazz);
        executor.clazz = clazz;
        executor.batchSize = batchSize;
        executor.listConsumer = listConsumer;
        return executor;
    }
    /* ---------- 自定义 class 映射 end ---------- */

    /* ---------- 自定义校验构造 ---------- */
    public static <T> EasyExcelBatchExecutor<T> builderValid(final Class<T> clazz, final BiConsumer<List<T>, Map<T, String>> listValidateConsumer) {
        EasyExcelBatchExecutor<T> executor = new EasyExcelBatchExecutor<>(DEFAULT_EXCEL_BATCH, false, clazz);
        executor.clazz = clazz;
        executor.listValidateConsumer = listValidateConsumer;
        return executor;
    }

    public static <T> EasyExcelBatchExecutor<T> builderValid(final int batchSize, final Class<T> clazz, final BiConsumer<List<T>, Map<T, String>> listValidateConsumer) {
        EasyExcelBatchExecutor<T> executor = new EasyExcelBatchExecutor<>(batchSize, false, clazz);
        executor.clazz = clazz;
        executor.batchSize = batchSize;
        executor.listValidateConsumer = listValidateConsumer;
        return executor;
    }
    /* ---------- 自定义校验 end ---------- */

    public EasyExcelBatchExecutor<T> excelHeaders(final List<String> excelHeaders) {
        this.excelHeaders = excelHeaders;
        return this;
    }

    public EasyExcelBatchExecutor<T> excelHeaders(final String[] excelHeaders) {
        this.excelHeaders = Arrays.asList(excelHeaders);
        return this;
    }

    public EasyExcelBatchExecutor<T> excelHeaderErrorMsg(final String excelHeaderErrorMsg) {
        this.excelHeaderErrorMsg = excelHeaderErrorMsg;
        return this;
    }

    public EasyExcelBatchExecutor<T> headerRowNumber(final int headerRowNumber) {
        this.headerRowNumber = headerRowNumber;
        return this;
    }

    public EasyExcelBatchExecutor<T> batchSize(final int batchSize) {
        this.batchSize = batchSize;
        return this;
    }

    public EasyExcelBatchExecutor<T> listConsumer(final BiConsumer<List<T>, Integer> listConsumer) {
        this.listConsumer = listConsumer;
        return this;
    }

    public EasyExcelBatchExecutor<T> setArrayList() {
        this.list = new ArrayList<>(this.list);
        return this;
    }

    public EasyExcelBatchExecutor<T> openClear() {
        this.clear = true;
        return this;
    }

    public EasyExcelBatchExecutor<T> nextClear(final boolean clearNextSheet) {
        this.clearNextSheet = clearNextSheet;
        return this;
    }

    public EasyExcelBatchExecutor<T> addSheetNo(Integer... sheetNo) {
        Arrays.stream(sheetNo).forEach(no -> this.sheetNoList.add(EasyExcel.readSheet(no)
                .head(clazz)
                .build()));
        return this;
    }

    public EasyExcelBatchExecutor<T> maxTotal(final int maxTotal) {
        this.maxTotal = maxTotal;
        return this;
    }

    public EasyExcelBatchExecutor<T> overflowTotalMsg(String overflowTotalMsg) {
        this.overflowTotalMsg = overflowTotalMsg;
        return this;
    }

    public List<T> getReadList() {
        return this.list;
    }

    public EasyExcelBatchExecutor<T> readBySax(InputStream in) {
        this.startRead(in);
        return this;
    }

    public EasyExcelBatchExecutor<T> readAllBySax(InputStream in) {
        this.readAll = true;
        this.startRead(in);
        return this;
    }

    public EasyExcelBatchExecutor<T> readBySax(InputStream in, int sheetNo) {
        this.sheetNo = sheetNo;
        this.startRead(in);
        return this;
    }

    public EasyExcelBatchExecutor<T> autoMappingHeaders() {
        ExcelHeader annotation = clazz.getAnnotation(ExcelHeader.class);
        if (annotation == null) {
            throw new EasyExcelException("映射类没配置 excelHeader 无法自动映射");
        }
        // 设置表头
        String[] value = annotation.value();
        this.excelHeaders = Arrays.asList(value);
        return this;
    }

    public EasyExcelBatchExecutor<T> validate(final boolean failedContinue, Function<T, String> failedFunction) {
        this.failedContinue = failedContinue;
        this.validateFunction = failedFunction;
        this.failedMap = new HashMap<>();
        return this;
    }

    public EasyExcelBatchExecutor<T> headerValidate(final Predicate<Map<Integer, String>> headerPredicate) {
        this.headerPredicate = headerPredicate;
        return this;
    }

    public EasyExcelBatchExecutor<T> hitCheckHeaderNum(Integer count) {
        this.hitCheckHeaderNum = count;
        return this;
    }

    public int getImportTotal() {
        return this.total;
    }

    @Override
    public void invokeHeadMap(Map<Integer, String> headMap, AnalysisContext context) {
        if (this.headerPredicate != null) {
            if (!headerPredicate.test(headMap)) {
                throw new EasyExcelException(Optional.ofNullable(this.excelHeaderErrorMsg).orElse("导入模板不正确"));
            }
        } else if (CollectionUtils.isNotEmpty(this.excelHeaders)) {
            this.recordHeaderCount++;
            Collection<String> headers = headMap.values();
            // 当前表头不满足需求表头
            if (headers.size() != this.excelHeaders.size() || !headers.containsAll(this.excelHeaders)) {
                if (this.hitCheckHeaderNum != null) {
                    // 只有满足命中行数 = 记录行数时才会抛出异常
                    if (this.hitCheckHeaderNum == this.recordHeaderCount) {
                        throw new EasyExcelException(Optional.ofNullable(this.excelHeaderErrorMsg).orElse("导入模板不正确"));
                    }
                } else {
                    throw new EasyExcelException(Optional.ofNullable(this.excelHeaderErrorMsg).orElse("导入模板不正确"));
                }
            }
        }
    }

    @Override
    public void invoke(T data, AnalysisContext context) {
        checkHeaders();

        this.total++;
        if (this.maxTotal > 0 && this.total > this.maxTotal) {
            throw new EasyExcelException(Optional.ofNullable(this.overflowTotalMsg).orElse(String.format("一次性导入不能超过 %d 条", this.maxTotal)));
        }

        // 用户自定义校验
        if (this.validateFunction != null) {
            // 校验失败
            String apply = this.validateFunction.apply(data);
            if (StringUtils.isNotBlank(apply)) {
                this.failedMap.put(data, apply);
                // 不放行就跳出
                if (!this.failedContinue) {
                    return;
                }
            }
        }

        // 清除分批，先 add 再处理
        if (this.clear) {
            executeClearBatch(data);
        } else {
            executeLogicBatch(data);
        }
    }

    @Override
    public void doAfterAllAnalysed(AnalysisContext context) {
        checkHeaders();

        // 如果是支持清除操作，这里可以进行再一次分批；最后结尾不进行清空
        if (this.clear) {
            if (this.list.size() >= this.batchSize) {
                List<List<T>> partition = Lists.partition(this.list, this.batchSize);
                for (List<T> tempList : partition) {
                    this.executeClear(tempList);
                }
            } else {
                this.executeClear(this.list);
            }
            // 翻页清空数据
            if (this.hasNextSheetNo() && this.clearNextSheet) {
                this.list.clear();
            }
        } else {
            // 处理剩余的数据
            if (logicBatch < list.size()) {
                this.executeLogicClear(this.list);
            }
            // 翻页清空数据
            if (this.hasNextSheetNo() && this.clearNextSheet) {
                this.list.clear();
            }
            // 初始化逻辑分批
            this.logicBatch = 0;
        }
        this.currentNo++;
    }

    @Override
    public void onException(Exception exception, AnalysisContext context) throws Exception {
        log.info("excel 解析第 {} 异常 -> {}:{}", context.readRowHolder().getRowIndex(),
                context.readRowHolder().getRowType(),
                context.readRowHolder().getCurrentRowAnalysisResult(),
                exception);
        // 抛出异常
        if (exception != null) {
            throw exception;
        }
    }

    /**
     * 当设置了表头之后，记录表头数要与实际指定表头数相同
     */
    private void checkHeaders() {
        if (CollectionUtils.isNotEmpty(this.excelHeaders) && (this.recordHeaderCount != this.headerRowNumber)) {
            throw new EasyExcelException(Optional.ofNullable(this.excelHeaderErrorMsg).orElse("导入模板不正确"));
        }
    }

    private void startRead(InputStream inputStream) {
        if (!inputStream.markSupported()) {
            inputStream = new BufferedInputStream(inputStream);
        }

        try {
            ExcelTypeEnum excelTypeEnum;
            FileMagic fileMagic = FileMagic.valueOf(inputStream);
            if (FileMagic.OLE2.equals(fileMagic)) {
                excelTypeEnum = ExcelTypeEnum.XLS;
            } else if (FileMagic.OOXML.equals(fileMagic)) {
                excelTypeEnum = ExcelTypeEnum.XLSX;
            } else {
                excelTypeEnum = ExcelTypeEnum.CSV;
            }
            ExcelReaderBuilder read = EasyExcelFactory.read(inputStream, clazz, this)
                    .excelType(excelTypeEnum)
                    .autoCloseStream(true)
                    .headRowNumber(this.headerRowNumber);
            if (this.readAll) {
                read.doReadAll();
            } else if (CollectionUtils.isNotEmpty(this.sheetNoList)) {
                read.build().read(this.sheetNoList).finish();
            } else {
                read.sheet(this.sheetNo).doRead();
            }
        } catch (IOException e) {
            throw new EasyExcelException(e);
        }
    }

    private boolean hasNextSheetNo() {
        if (CollectionUtils.isNotEmpty(this.sheetNoList)) {
            return this.sheetNoList.size() > 1;
        }
        return false;
    }

    /**
     * 批量处理清除分批
     *
     * @param data 每次处理的数据
     */
    private void executeClearBatch(T data) {
        this.list.add(data);
        if (this.list.size() >= this.batchSize) {
            this.executeClear(this.list);
        }
    }

    /**
     * 批量处理逻辑分批
     *
     * @param data 每次处理的数据
     */
    private void executeLogicBatch(T data) {
        // 逻辑分批先判断再 add
        if (this.total > 0 && this.total % this.batchSize == 0) {
            this.executeLogicClear(this.list);
            this.logicBatch = list.size();
        }
        this.list.add(data);
    }

    private void executeClear(List<T> executeList) {
        if (CollectionUtils.isEmpty(executeList)) {
            return;
        }

        if (this.listValidateConsumer != null) {
            this.listValidateConsumer.andThen((tmpList, tmpCurNo) -> tmpList.clear()).accept(executeList, this.failedMap);
        } else if (this.listConsumer != null) {
            this.listConsumer.andThen((tmpList, tmpCurNo) -> tmpList.clear()).accept(executeList, this.currentNo);
        }
    }

    private void executeLogicClear(List<T> executeList) {
        if (CollectionUtils.isEmpty(executeList)) {
            return;
        }

        if (this.listValidateConsumer != null) {
            this.listValidateConsumer.accept(executeList.subList(logicBatch, executeList.size()), this.failedMap);
        } else if (this.listConsumer != null) {
            this.listConsumer.accept(executeList.subList(logicBatch, executeList.size()), this.currentNo);
        }
    }
}
